{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SQL Auto Vector Query Engine\n",
    "In this tutorial, we show you how to use our SQLAutoVectorQueryEngine.\n",
    "\n",
    "This query engine allows you to combine insights from your structured tables with your unstructured data.\n",
    "It first decides whether to query your structured tables for insights.\n",
    "Once it does, it can then infer a corresponding query to the vector store in order to fetch corresponding documents."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openai\n",
    "import os\n",
    "\n",
    "os.environ[\"OPENAI_API_KEY\"] = \"[You API key]\"\n",
    "openai.api_key = os.environ[\"OPENAI_API_KEY\"]"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# NOTE: This is ONLY necessary in jupyter notebook.\n",
    "# Details: Jupyter runs an event-loop behind the scenes.\n",
    "#          This results in nested event-loops when we start an event-loop to make async queries.\n",
    "#          This is normally not allowed, we use nest_asyncio to allow it for convenience.\n",
    "import nest_asyncio\n",
    "\n",
    "nest_asyncio.apply()\n",
    "\n",
    "import logging\n",
    "import sys\n",
    "\n",
    "logging.basicConfig(stream=sys.stdout, level=logging.INFO)\n",
    "logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "INFO:numexpr.utils:Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "Note: NumExpr detected 12 cores but \"NUMEXPR_MAX_THREADS\" not set, so enforcing safe limit of 8.\n",
      "INFO:numexpr.utils:NumExpr defaulting to 8 threads.\n",
      "NumExpr defaulting to 8 threads.\n"
     ]
    }
   ],
   "source": [
    "from llama_index import (\n",
    "    VectorStoreIndex,\n",
    "    SimpleDirectoryReader,\n",
    "    ServiceContext,\n",
    "    StorageContext,\n",
    "    SQLDatabase,\n",
    "    WikipediaReader,\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Common Objects\n",
    "\n",
    "This includes a `ServiceContext` object containing abstractions such as the LLM and chunk size.\n",
    "This also includes a `StorageContext` object containing our vector store abstractions."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages/pinecone/index.py:4: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from tqdm.autonotebook import tqdm\n"
     ]
    }
   ],
   "source": [
    "# define pinecone index\n",
    "import pinecone\n",
    "import os\n",
    "\n",
    "api_key = os.environ[\"PINECONE_API_KEY\"]\n",
    "pinecone.init(api_key=api_key, environment=\"us-west1-gcp-free\")\n",
    "\n",
    "# dimensions are for text-embedding-ada-002\n",
    "# pinecone.create_index(\"quickstart\", dimension=1536, metric=\"euclidean\", pod_type=\"p1\")\n",
    "pinecone_index = pinecone.Index(\"quickstart\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "{}"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# OPTIONAL: delete all\n",
    "pinecone_index.delete(deleteAll=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.node_parser.simple import SimpleNodeParser\n",
    "from llama_index import ServiceContext, LLMPredictor\n",
    "from llama_index.storage import StorageContext\n",
    "from llama_index.vector_stores import PineconeVectorStore\n",
    "from llama_index.text_splitter import TokenTextSplitter\n",
    "from llama_index.llms import OpenAI\n",
    "\n",
    "# define node parser and LLM\n",
    "chunk_size = 1024\n",
    "llm = OpenAI(temperature=0, model=\"gpt-4\", streaming=True)\n",
    "service_context = ServiceContext.from_defaults(chunk_size=chunk_size, llm=llm)\n",
    "text_splitter = TokenTextSplitter(chunk_size=chunk_size)\n",
    "node_parser = SimpleNodeParser.from_defaults(text_splitter=text_splitter)\n",
    "\n",
    "# define pinecone vector index\n",
    "vector_store = PineconeVectorStore(\n",
    "    pinecone_index=pinecone_index, namespace=\"wiki_cities\"\n",
    ")\n",
    "storage_context = StorageContext.from_defaults(vector_store=vector_store)\n",
    "vector_index = VectorStoreIndex([], storage_context=storage_context)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Create Database Schema + Test Data\n",
    "\n",
    "Here we introduce a toy scenario where there are 100 tables (too big to fit into the prompt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sqlalchemy import (\n",
    "    create_engine,\n",
    "    MetaData,\n",
    "    Table,\n",
    "    Column,\n",
    "    String,\n",
    "    Integer,\n",
    "    select,\n",
    "    column,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "engine = create_engine(\"sqlite:///:memory:\", future=True)\n",
    "metadata_obj = MetaData()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# create city SQL table\n",
    "table_name = \"city_stats\"\n",
    "city_stats_table = Table(\n",
    "    table_name,\n",
    "    metadata_obj,\n",
    "    Column(\"city_name\", String(16), primary_key=True),\n",
    "    Column(\"population\", Integer),\n",
    "    Column(\"country\", String(16), nullable=False),\n",
    ")\n",
    "\n",
    "metadata_obj.create_all(engine)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "dict_keys(['city_stats'])"
      ]
     },
     "execution_count": 10,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# print tables\n",
    "metadata_obj.tables.keys()"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "We introduce some test data into the `city_stats` table"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from sqlalchemy import insert\n",
    "\n",
    "rows = [\n",
    "    {\"city_name\": \"Toronto\", \"population\": 2930000, \"country\": \"Canada\"},\n",
    "    {\"city_name\": \"Tokyo\", \"population\": 13960000, \"country\": \"Japan\"},\n",
    "    {\"city_name\": \"Berlin\", \"population\": 3645000, \"country\": \"Germany\"},\n",
    "]\n",
    "for row in rows:\n",
    "    stmt = insert(city_stats_table).values(**row)\n",
    "    with engine.connect() as connection:\n",
    "        cursor = connection.execute(stmt)\n",
    "        connection.commit()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[('Toronto', 2930000, 'Canada'), ('Tokyo', 13960000, 'Japan'), ('Berlin', 3645000, 'Germany')]\n"
     ]
    }
   ],
   "source": [
    "with engine.connect() as connection:\n",
    "    cursor = connection.exec_driver_sql(\"SELECT * FROM city_stats\")\n",
    "    print(cursor.fetchall())"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Load Data\n",
    "\n",
    "We first show how to convert a Document into a set of Nodes, and insert into a DocumentStore."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Requirement already satisfied: wikipedia in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (1.4.0)\n",
      "Requirement already satisfied: beautifulsoup4 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from wikipedia) (4.12.2)\n",
      "Requirement already satisfied: requests<3.0.0,>=2.0.0 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from wikipedia) (2.31.0)\n",
      "Requirement already satisfied: idna<4,>=2.5 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (3.4)\n",
      "Requirement already satisfied: charset-normalizer<4,>=2 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (3.2.0)\n",
      "Requirement already satisfied: certifi>=2017.4.17 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (2023.5.7)\n",
      "Requirement already satisfied: urllib3<3,>=1.21.1 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from requests<3.0.0,>=2.0.0->wikipedia) (1.26.16)\n",
      "Requirement already satisfied: soupsieve>1.2 in /Users/loganmarkewich/llama_index/llama-index/lib/python3.9/site-packages (from beautifulsoup4->wikipedia) (2.4.1)\n",
      "\u001b[33mWARNING: You are using pip version 21.2.4; however, version 23.2 is available.\n",
      "You should consider upgrading via the '/Users/loganmarkewich/llama_index/llama-index/bin/python3 -m pip install --upgrade pip' command.\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "# install wikipedia python package\n",
    "!pip install wikipedia"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "cities = [\"Toronto\", \"Berlin\", \"Tokyo\"]\n",
    "wiki_docs = WikipediaReader().load_data(pages=cities)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build SQL Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql_database = SQLDatabase(engine, include_tables=[\"city_stats\"])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "from llama_index.indices.struct_store.sql_query import NLSQLTableQueryEngine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "sql_query_engine = NLSQLTableQueryEngine(\n",
    "    sql_database=sql_database,\n",
    "    tables=[\"city_stats\"],\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Build Vector Index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "Upserted vectors: 100%|██████████| 20/20 [00:00<00:00, 22.37it/s]\n",
      "Upserted vectors: 100%|██████████| 22/22 [00:00<00:00, 23.14it/s]\n",
      "Upserted vectors: 100%|██████████| 13/13 [00:00<00:00, 17.67it/s]\n"
     ]
    }
   ],
   "source": [
    "# Insert documents into vector index\n",
    "# Each document has metadata of the city attached\n",
    "for city, wiki_doc in zip(cities, wiki_docs):\n",
    "    nodes = node_parser.get_nodes_from_documents([wiki_doc])\n",
    "    # add metadata to each node\n",
    "    for node in nodes:\n",
    "        node.metadata = {\"title\": city}\n",
    "    vector_index.insert_nodes(nodes)"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define Query Engines, Set as Tools"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.query_engine import SQLAutoVectorQueryEngine, RetrieverQueryEngine\n",
    "from llama_index.tools.query_engine import QueryEngineTool\n",
    "from llama_index.indices.vector_store import VectorIndexAutoRetriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from llama_index.indices.vector_store.retrievers import VectorIndexAutoRetriever\n",
    "from llama_index.vector_stores.types import MetadataInfo, VectorStoreInfo\n",
    "from llama_index.query_engine.retriever_query_engine import RetrieverQueryEngine\n",
    "\n",
    "\n",
    "vector_store_info = VectorStoreInfo(\n",
    "    content_info=\"articles about different cities\",\n",
    "    metadata_info=[\n",
    "        MetadataInfo(name=\"title\", type=\"str\", description=\"The name of the city\"),\n",
    "    ],\n",
    ")\n",
    "vector_auto_retriever = VectorIndexAutoRetriever(\n",
    "    vector_index, vector_store_info=vector_store_info\n",
    ")\n",
    "\n",
    "retriever_query_engine = RetrieverQueryEngine.from_args(\n",
    "    vector_auto_retriever, service_context=service_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "sql_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=sql_query_engine,\n",
    "    description=(\n",
    "        \"Useful for translating a natural language query into a SQL query over a table containing: \"\n",
    "        \"city_stats, containing the population/country of each city\"\n",
    "    ),\n",
    ")\n",
    "vector_tool = QueryEngineTool.from_defaults(\n",
    "    query_engine=retriever_query_engine,\n",
    "    description=f\"Useful for answering semantic questions about different cities\",\n",
    ")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Define SQLAutoVectorQueryEngine"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "query_engine = SQLAutoVectorQueryEngine(\n",
    "    sql_tool, vector_tool, service_context=service_context\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 23,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3mQuerying SQL database: Useful for translating a natural language query into a SQL query over a table containing city_stats, containing the population/country of each city\n",
      "\u001b[0mINFO:llama_index.query_engine.sql_join_query_engine:> Querying SQL database: Useful for translating a natural language query into a SQL query over a table containing city_stats, containing the population/country of each city\n",
      "> Querying SQL database: Useful for translating a natural language query into a SQL query over a table containing city_stats, containing the population/country of each city\n",
      "INFO:llama_index.indices.struct_store.sql_query:> Table desc str: Table 'city_stats' has columns: city_name (VARCHAR(16)), population (INTEGER), country (VARCHAR(16)) and foreign keys: .\n",
      "> Table desc str: Table 'city_stats' has columns: city_name (VARCHAR(16)), population (INTEGER), country (VARCHAR(16)) and foreign keys: .\n",
      "\u001b[33;1m\u001b[1;3mSQL query: SELECT city_name, population FROM city_stats ORDER BY population DESC LIMIT 1;\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3mSQL response: \n",
      "Tokyo is the city with the highest population, with 13.96 million people. It is a vibrant city with a rich culture and a wide variety of art forms. From traditional Japanese art such as calligraphy and woodblock prints to modern art galleries and museums, Tokyo has something for everyone. There are also many festivals and events throughout the year that celebrate the city's culture and art.\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mTransformed query given SQL response: What are some specific cultural festivals, events, and notable art galleries or museums in Tokyo?\n",
      "\u001b[0mINFO:llama_index.query_engine.sql_join_query_engine:> Transformed query given SQL response: What are some specific cultural festivals, events, and notable art galleries or museums in Tokyo?\n",
      "> Transformed query given SQL response: What are some specific cultural festivals, events, and notable art galleries or museums in Tokyo?\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using query str: cultural festivals events art galleries museums Tokyo\n",
      "Using query str: cultural festivals events art galleries museums Tokyo\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using filters: {'title': 'Tokyo'}\n",
      "Using filters: {'title': 'Tokyo'}\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using top_k: 2\n",
      "Using top_k: 2\n",
      "\u001b[38;5;200m\u001b[1;3mquery engine response: The context information mentions the Tokyo National Museum, which houses 37% of the country's artwork national treasures. It also mentions the Studio Ghibli anime center as a subcultural attraction. However, the text does not provide information on specific cultural festivals or events in Tokyo.\n",
      "\u001b[0mINFO:llama_index.query_engine.sql_join_query_engine:> query engine response: The context information mentions the Tokyo National Museum, which houses 37% of the country's artwork national treasures. It also mentions the Studio Ghibli anime center as a subcultural attraction. However, the text does not provide information on specific cultural festivals or events in Tokyo.\n",
      "> query engine response: The context information mentions the Tokyo National Museum, which houses 37% of the country's artwork national treasures. It also mentions the Studio Ghibli anime center as a subcultural attraction. However, the text does not provide information on specific cultural festivals or events in Tokyo.\n",
      "\u001b[32;1m\u001b[1;3mFinal response: Tokyo, the city with the highest population of 13.96 million people, is known for its vibrant culture and diverse art forms. It is home to traditional Japanese art such as calligraphy and woodblock prints, as well as modern art galleries and museums. Notably, the Tokyo National Museum houses 37% of the country's artwork national treasures, and the Studio Ghibli anime center is a popular subcultural attraction. While there are many festivals and events throughout the year that celebrate the city's culture and art, specific examples were not provided in the available information.\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\n",
    "    \"Tell me about the arts and culture of the city with the highest population\"\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 24,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Tokyo, the city with the highest population of 13.96 million people, is known for its vibrant culture and diverse art forms. It is home to traditional Japanese art such as calligraphy and woodblock prints, as well as modern art galleries and museums. Notably, the Tokyo National Museum houses 37% of the country's artwork national treasures, and the Studio Ghibli anime center is a popular subcultural attraction. While there are many festivals and events throughout the year that celebrate the city's culture and art, specific examples were not provided in the available information.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 25,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3mQuerying other query engine: Useful for answering semantic questions about different cities\n",
      "\u001b[0mINFO:llama_index.query_engine.sql_join_query_engine:> Querying other query engine: Useful for answering semantic questions about different cities\n",
      "> Querying other query engine: Useful for answering semantic questions about different cities\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using query str: history of Berlin\n",
      "Using query str: history of Berlin\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using filters: {'title': 'Berlin'}\n",
      "Using filters: {'title': 'Berlin'}\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using top_k: 2\n",
      "Using top_k: 2\n",
      "\u001b[38;5;200m\u001b[1;3mQuery Engine response: Berlin's history dates back to around 60,000 BC, with the earliest human traces found in the area. A Mesolithic deer antler mask found in Biesdorf (Berlin) was dated around 9000 BC. During Neolithic times, a large number of communities existed in the area and in the Bronze Age, up to 1000 people lived in 50 villages. Early Germanic tribes took settlement from 500 BC and Slavic settlements and castles began around 750 AD.\n",
      "\n",
      "The earliest evidence of middle age settlements in the area of today's Berlin are remnants of a house foundation dated to 1174, found in excavations in Berlin Mitte, and a wooden beam dated from approximately 1192. The first written records of towns in the area of present-day Berlin date from the late 12th century. Spandau is first mentioned in 1197 and Köpenick in 1209, although these areas did not join Berlin until 1920. \n",
      "\n",
      "The central part of Berlin can be traced back to two towns. Cölln on the Fischerinsel is first mentioned in a 1237 document, and Berlin, across the Spree in what is now called the Nikolaiviertel, is referenced in a document from 1244. 1237 is considered the founding date of the city. The two towns over time formed close economic and social ties, and profited from the staple right on the two important trade routes Via Imperii and from Bruges to Novgorod. In 1307, they formed an alliance with a common external policy, their internal administrations still being separated. In 1415, Frederick I became the elector of the Margraviate of Brandenburg, which he ruled until 1440.\n",
      "\n",
      "The name Berlin has its roots in the language of West Slavic inhabitants of the area of today's Berlin, and may be related to the Old Polabian stem berl-/birl- (\"swamp\"). or Proto-Slavic bьrlogъ, (lair, den). Since the Ber- at the beginning sounds like the German word Bär (\"bear\"), a bear appears in the coat of arms of the city. It is therefore an example of canting arms.\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"Tell me about the history of Berlin\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 26,
   "metadata": {
    "scrolled": true,
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Berlin's history dates back to around 60,000 BC, with the earliest human traces found in the area. A Mesolithic deer antler mask found in Biesdorf (Berlin) was dated around 9000 BC. During Neolithic times, a large number of communities existed in the area and in the Bronze Age, up to 1000 people lived in 50 villages. Early Germanic tribes took settlement from 500 BC and Slavic settlements and castles began around 750 AD.\n",
      "\n",
      "The earliest evidence of middle age settlements in the area of today's Berlin are remnants of a house foundation dated to 1174, found in excavations in Berlin Mitte, and a wooden beam dated from approximately 1192. The first written records of towns in the area of present-day Berlin date from the late 12th century. Spandau is first mentioned in 1197 and Köpenick in 1209, although these areas did not join Berlin until 1920. \n",
      "\n",
      "The central part of Berlin can be traced back to two towns. Cölln on the Fischerinsel is first mentioned in a 1237 document, and Berlin, across the Spree in what is now called the Nikolaiviertel, is referenced in a document from 1244. 1237 is considered the founding date of the city. The two towns over time formed close economic and social ties, and profited from the staple right on the two important trade routes Via Imperii and from Bruges to Novgorod. In 1307, they formed an alliance with a common external policy, their internal administrations still being separated. In 1415, Frederick I became the elector of the Margraviate of Brandenburg, which he ruled until 1440.\n",
      "\n",
      "The name Berlin has its roots in the language of West Slavic inhabitants of the area of today's Berlin, and may be related to the Old Polabian stem berl-/birl- (\"swamp\"). or Proto-Slavic bьrlogъ, (lair, den). Since the Ber- at the beginning sounds like the German word Bär (\"bear\"), a bear appears in the coat of arms of the city. It is therefore an example of canting arms.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 27,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[36;1m\u001b[1;3mQuerying SQL database: Useful for translating a natural language query into a SQL query over a table containing: city_stats, containing the population/country of each city\n",
      "\u001b[0mINFO:llama_index.query_engine.sql_join_query_engine:> Querying SQL database: Useful for translating a natural language query into a SQL query over a table containing: city_stats, containing the population/country of each city\n",
      "> Querying SQL database: Useful for translating a natural language query into a SQL query over a table containing: city_stats, containing the population/country of each city\n",
      "INFO:llama_index.indices.struct_store.sql_query:> Table desc str: Table 'city_stats' has columns: city_name (VARCHAR(16)), population (INTEGER), country (VARCHAR(16)) and foreign keys: .\n",
      "> Table desc str: Table 'city_stats' has columns: city_name (VARCHAR(16)), population (INTEGER), country (VARCHAR(16)) and foreign keys: .\n",
      "\u001b[33;1m\u001b[1;3mSQL query: SELECT city_name, country FROM city_stats;\n",
      "\u001b[0m\u001b[33;1m\u001b[1;3mSQL response:  Toronto is in Canada, Tokyo is in Japan, and Berlin is in Germany.\n",
      "\u001b[0m\u001b[36;1m\u001b[1;3mTransformed query given SQL response: What countries are New York, San Francisco, and other cities in?\n",
      "\u001b[0mINFO:llama_index.query_engine.sql_join_query_engine:> Transformed query given SQL response: What countries are New York, San Francisco, and other cities in?\n",
      "> Transformed query given SQL response: What countries are New York, San Francisco, and other cities in?\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using query str: New York San Francisco\n",
      "Using query str: New York San Francisco\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using filters: {'title': 'San Francisco'}\n",
      "Using filters: {'title': 'San Francisco'}\n",
      "INFO:llama_index.indices.vector_store.retrievers.auto_retriever.auto_retriever:Using top_k: 2\n",
      "Using top_k: 2\n",
      "\u001b[38;5;200m\u001b[1;3mquery engine response: None\n",
      "\u001b[0mINFO:llama_index.query_engine.sql_join_query_engine:> query engine response: None\n",
      "> query engine response: None\n",
      "\u001b[32;1m\u001b[1;3mFinal response: The country corresponding to each city is as follows: Toronto is in Canada, Tokyo is in Japan, and Berlin is in Germany. Unfortunately, I do not have information on the countries for New York, San Francisco, and other cities.\n",
      "\u001b[0m"
     ]
    }
   ],
   "source": [
    "response = query_engine.query(\"Can you give me the country corresponding to each city?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 28,
   "metadata": {
    "tags": []
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "The country corresponding to each city is as follows: Toronto is in Canada, Tokyo is in Japan, and Berlin is in Germany. Unfortunately, I do not have information on the countries for New York, San Francisco, and other cities.\n"
     ]
    }
   ],
   "source": [
    "print(str(response))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "llama-index",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
